package com.rocpenghua.excel.easyexcel;

import cn.hutool.core.util.ReflectUtil;
import com.alibaba.excel.write.handler.SheetWriteHandler;
import com.alibaba.excel.write.metadata.holder.WriteSheetHolder;
import com.alibaba.excel.write.metadata.holder.WriteWorkbookHolder;
import com.rocpenghua.excel.water.ImageParam;
import com.rocpenghua.excel.water.ImageUtils;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.apache.poi.openxml4j.opc.PackagePartName;
import org.apache.poi.openxml4j.opc.PackageRelationship;
import org.apache.poi.openxml4j.opc.TargetMode;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.streaming.SXSSFSheet;
import org.apache.poi.xssf.streaming.SXSSFWorkbook;
import org.apache.poi.xssf.usermodel.XSSFPictureData;
import org.apache.poi.xssf.usermodel.XSSFRelation;
import org.apache.poi.xssf.usermodel.XSSFSheet;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;

/**
 * @author pengp
 * @date 2024年02月27日 11:27
 */
@RequiredArgsConstructor
public class WaterMarkHandler implements SheetWriteHandler {

    private final String WATER_MARK;

    public static ByteArrayOutputStream createWaterMark(String content) throws IOException {
        List<String> contents = Arrays.asList(content.split(","));
        final BufferedImage image = ImageUtils.createWaterMark(contents, new ImageParam());
        ByteArrayOutputStream os = new ByteArrayOutputStream();
        ImageIO.write(image, "png", os);
        return os;
    }

    /**

     为Excel打上水印工具函数
     @param sheet excel sheet
     @param bytes 水印图片字节数组
     */
    public static void putWaterRemarkToExcel(SXSSFSheet sheet, byte[] bytes) {
        //add relation from sheet to the picture data
        SXSSFWorkbook workbook = sheet.getWorkbook();
        XSSFSheet shReflect = (XSSFSheet) ReflectUtil.getFieldValue(sheet, "_sh");
        int pictureIdx = workbook.addPicture(bytes, Workbook.PICTURE_TYPE_PNG);
        XSSFPictureData xssfPictureData = (XSSFPictureData) workbook.getAllPictures().get(pictureIdx);
        PackagePartName ppn = xssfPictureData.getPackagePart().getPartName();
        PackageRelationship pr = shReflect.getPackagePart().addRelationship(ppn, TargetMode.INTERNAL, XSSFRelation.IMAGES.getRelation(), null);
        //set background picture to sheet
        shReflect.getCTWorksheet().addNewPicture().setId(pr.getId());
    }
    @Override
    public void beforeSheetCreate(WriteWorkbookHolder writeWorkbookHolder, WriteSheetHolder writeSheetHolder) {
    }

    @SneakyThrows
    @Override
    public void afterSheetCreate(WriteWorkbookHolder writeWorkbookHolder, WriteSheetHolder writeSheetHolder) {
        try (ByteArrayOutputStream waterMark = createWaterMark(WATER_MARK)){
            SXSSFSheet sheet = (SXSSFSheet) writeSheetHolder.getSheet();
            putWaterRemarkToExcel(sheet, waterMark.toByteArray());
        }
    }
}
